/*
 * Copyright (c) 2024 Huawei Device Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "keystore.h"

#include <asn1/asn1.h>
#include <asn1/oid.h>
#include <credentials/builder.h>
#include <credentials/certificates/x509.h>
#include <credentials/keys/public_key.h>
#include <credentials/sets/mem_cred.h>
#include <encoding/payloads/auth_payload.h>
#include <encoding/payloads/hash_payload.h>

#include <plugins/stroke/stroke_ca.h>
#include <sa/ikev2/keymat_v2.h>
#include <sa/ikev1/keymat_v1.h>
#include <securec.h>
#include <string.h>

#include <cert_manager_api.h>

// copy from openssl to avoid asn1 duplicate defined
#define RSA_PKCS1_PADDING_SIZE 11
// RSA_size(rsa)
#define RSA_SIZE 256

int keystore_load_ca_cert(stroke_ca_t *ca, mem_cred_t *creds)
{
	DBG1(DBG_CFG, "loading ca certificate from keystore");
	chunk_t ca_blob = load_vpn_cert(CA_CERT);

	if (!ca_blob.ptr)
	{
		DBG1(DBG_CFG, "loading ca certificate from keystore error");
		return -1;
	}

	certificate_t *cert = NULL;
	cert = lib->creds->create(lib->creds, CRED_CERTIFICATE, CERT_X509, BUILD_BLOB_PEM, ca_blob, BUILD_X509_FLAG,
							  X509_CA, BUILD_END);
	if (!cert)
	{
		cert = lib->creds->create(lib->creds, CRED_CERTIFICATE, CERT_X509, BUILD_BLOB_ASN1_DER, ca_blob,
								  BUILD_X509_FLAG, X509_CA, BUILD_END);
	}
	chunk_free(&ca_blob);
	if (cert)
	{
		x509_t *x509 = (x509_t *)cert;

		if (!(x509->get_flags(x509) & X509_CA))
		{
			DBG1(DBG_CFG,
				 "  ca certificate \"%Y\" lacks ca basic constraint, "
				 "discarded",
				 cert->get_subject(cert));
			cert->destroy(cert);
			return -1;
		}
		DBG1(DBG_CFG, "  loaded ca certificate \"%Y\" from keystore", cert->get_subject(cert));

		cert = ca->get_cert_ref(ca, cert);
		creds->add_cert(creds, TRUE, cert);
		return 0;
	}
	else
	{
		DBG1(DBG_CFG, "  loading ca certificate from keystore failed");
		return -1;
	}
}

certificate_t *keystore_load_user_cert()
{
	DBG1(DBG_CFG, "loading user certificate from keystore");
	chunk_t user_blob = load_vpn_cert(USER_CERT);
	if (!user_blob.ptr)
	{
		DBG1(DBG_CFG, "loading user certificate from keystore error");
		return NULL;
	}

	certificate_t *cert = NULL;
	cert = lib->creds->create(lib->creds, CRED_CERTIFICATE, CERT_X509, BUILD_BLOB_PEM, user_blob, BUILD_END);
	if (!cert)
	{
		cert = lib->creds->create(lib->creds, CRED_CERTIFICATE, CERT_X509, BUILD_BLOB_ASN1_DER, user_blob, BUILD_END);
	}
	if (cert)
	{
		DBG1(DBG_CFG, "  loaded user certificate \"%Y\" from keystore", cert->get_subject(cert));
	}
	else
	{
		DBG1(DBG_CFG, "  loaded user certificate from keystore failed. ");
	}
	chunk_free(&user_blob);
	return cert;
}

static bool keystore_signature_params_build(signature_scheme_t scheme, chunk_t *asn1)
{
	int oid;

	oid = signature_scheme_to_oid(scheme);
	DBG1(DBG_CFG, "keystore_signature_params_build oid = %d", oid);
	if (oid == OID_UNKNOWN)
	{
		return FALSE;
	}
	DBG1(DBG_CFG, "keystore_signature_params_build asn1_algorithmIdentifier start");
	*asn1 = asn1_algorithmIdentifier(oid);
	return TRUE;
}

static bool keystore_build_signature_auth_data(chunk_t *auth_data, signature_scheme_t scheme)
{
	chunk_t data;
	uint8_t len;
	if (!keystore_signature_params_build(scheme, &data))
	{
		chunk_free(auth_data);
		return FALSE;
	}
	len = data.len;
	*auth_data = chunk_cat("cmm", chunk_from_thing(len), data, *auth_data);
	return TRUE;
}

static void keystore_add_auth_to_message(message_t *message, auth_method_t method, chunk_t data, bool notify)
{
	auth_payload_t *auth_payload;

	if (notify)
	{
		message->add_notify(message, FALSE, NO_PPK_AUTH, data);
	}
	else
	{
		auth_payload = auth_payload_create();
		auth_payload->set_auth_method(auth_payload, method);
		auth_payload->set_data(auth_payload, data);
		message->add_payload(message, (payload_t *)auth_payload);
	}
	chunk_free(&data);
}

// IKEv2 message sign
status_t keystore_ikev2_sign(ike_sa_t *ike_sa, chunk_t ike_sa_init, char *reserved, chunk_t ppk, chunk_t nonce,
							 message_t *message, identification_t *id)
{
	DBG1(DBG_IKE, " keystore_ikev2_sign start...");
	if (!ike_sa || !reserved)
	{
		return FAILED;
	}

	chunk_t octets = chunk_empty;
	keymat_v2_t *keymat = (keymat_v2_t *)ike_sa->get_keymat(ike_sa);
	keymat->get_auth_octets(keymat, FALSE, ike_sa_init, nonce, ppk, id, reserved, &octets, NULL);
	chunk_t auth_data;

	uint8_t uri_buf[MAX_LEN_URI];
	memset_s(uri_buf, MAX_LEN_URI, 0, MAX_LEN_URI);
	uint32_t uri_size = LoadVpnCertUri(USER_CERT, uri_buf, MAX_LEN_URI);
	if (uri_size < 1)
	{
		DBG1(DBG_IKE, "keystore_ikev2_sign cm_sign failed: user cert not found");
		return NOT_FOUND;
	}

	struct CmSignatureSpec signSpec = {CM_KEY_PURPOSE_SIGN, CM_PADDING_PKCS1_V1_5, CM_DIGEST_SHA256};
	int32_t result = cm_sign(uri_buf, MAX_LEN_URI, octets, &auth_data, &signSpec);
	if (result != 0)
	{
		DBG1(DBG_IKE, "keystore_ikev2_sign cm_sign failed");
		return NOT_FOUND;
	}
	signature_scheme_t scheme = SIGN_RSA_EMSA_PKCS1_SHA2_256;
	if (!keystore_build_signature_auth_data(&auth_data, scheme))
	{
		DBG1(DBG_IKE, "keystore_ikev2_sign keystore_build_signature_auth_data failed");
		return NOT_FOUND;
	}

	keystore_add_auth_to_message(message, AUTH_DS, auth_data, FALSE);
	chunk_free(&octets);
	DBG1(DBG_IKE, "keystore_ikev2_sign cm_sign success");
	return SUCCESS;
}

static int RSA_padding_add_PKCS1_type_1(unsigned char *to, int tlen,
										const unsigned char *from, int flen)
{
	DBG1(DBG_IKE, "RSA_padding_add_PKCS1_type_1");
	int j;
	unsigned char *p;

	if (flen > (tlen - RSA_PKCS1_PADDING_SIZE))
	{
		DBG1(DBG_IKE, "RSA_padding_add_PKCS1_type_1: RSA_R_DATA_TOO_LARGE_FOR_KEY_SIZE");
		return 0;
	}

	p = (unsigned char *)to;

	*(p++) = 0;
	*(p++) = 1; /* Private Key BT (Block Type) */

	/* pad out with 0xff data */
	j = tlen - 3 - flen;
	memset(p, 0xff, j);
	p += j;
	*(p++) = '\0';
	memcpy(p, from, (unsigned int)flen);
	return 1;
}

// IKEv1 message sign
status_t keystore_ikev1_sign(ike_sa_t *ike_sa, bool initiator, key_exchange_t *dhke, chunk_t dh_value,
							 chunk_t sa_payload, chunk_t id_payload, message_t *message, identification_t *id)
{
	DBG1(DBG_IKE, "keystore_ikev1_sign start...");
	if (!ike_sa || !dhke)
	{
		return FAILED;
	}

	chunk_t hash, sig, dh;
	hash_payload_t *sig_payload;
	if (!dhke->get_public_key(dhke, &dh))
	{
		DBG1(DBG_IKE, "keystore_ikev1_sign get_public_key failed");
		return FAILED;
	}
	keymat_v1_t *keymat = (keymat_v1_t *)ike_sa->get_keymat(ike_sa);
	signature_scheme_t scheme = SIGN_RSA_EMSA_PKCS1_NULL;
	if (!keymat->get_hash(keymat, initiator, dh, dh_value, ike_sa->get_id(ike_sa), sa_payload, id_payload, &hash,
						  &scheme))
	{
		free(dh.ptr);
		DBG1(DBG_IKE, "keystore_ikev1_sign get_hash failed");
		return FAILED;
	}
	free(dh.ptr);
	DBG1(DBG_IKE, "keystore_ikev1_sign , scheme = %N", signature_scheme_names, scheme);
	uint8_t uri_buf[MAX_LEN_URI];
	memset_s(uri_buf, MAX_LEN_URI, 0, MAX_LEN_URI);
	uint32_t uri_size = LoadVpnCertUri(USER_CERT, uri_buf, MAX_LEN_URI);
	if (uri_size < 1)
	{
		DBG1(DBG_IKE, "keystore_ikev1_sign failed: user cert not found");
		return NOT_FOUND;
	}

	chunk_t padding = chunk_alloc(RSA_SIZE);
	if(!RSA_padding_add_PKCS1_type_1(padding.ptr, padding.len, hash.ptr, hash.len)) {
		return NOT_FOUND;
	}

	struct CmSignatureSpec signSpec = {CM_KEY_PURPOSE_SIGN, CM_PADDING_NONE, CM_DIGEST_NONE};
	int32_t result = cm_sign(uri_buf, MAX_LEN_URI, padding, &sig, &signSpec);
	if (result != 0)
	{
		DBG1(DBG_IKE, "keystore_ikev1_sign cm_sign failed");
		return NOT_FOUND;
	}

	sig_payload = hash_payload_create(PLV1_SIGNATURE);
	sig_payload->set_hash(sig_payload, sig);
	free(sig.ptr);
	message->add_payload(message, &sig_payload->payload_interface);

	DBG1(DBG_IKE, "keystore_ikev1_sign cm_sign success");
	return SUCCESS;
}
